package resource

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/louloulin/dataflare/pkg/cluster"
	"github.com/louloulin/dataflare/pkg/cluster/coordination"
)

// ResourceType 表示资源类型
type ResourceType string

const (
	// ResourceTypeCPU 表示CPU资源
	ResourceTypeCPU ResourceType = "cpu"
	// ResourceTypeMemory 表示内存资源
	ResourceTypeMemory ResourceType = "memory"
	// ResourceTypeDisk 表示磁盘资源
	ResourceTypeDisk ResourceType = "disk"
	// ResourceTypeNetwork 表示网络资源
	ResourceTypeNetwork ResourceType = "network"
	// ResourceTypeGPU 表示GPU资源
	ResourceTypeGPU ResourceType = "gpu"
)

// ResourceQuantity 表示资源数量
type ResourceQuantity struct {
	// Value 是资源值
	Value float64
	// Unit 是资源单位
	Unit string
}

// String 返回资源数量的字符串表示
func (q ResourceQuantity) String() string {
	return fmt.Sprintf("%g%s", q.Value, q.Unit)
}

// ResourceRequirement 表示资源需求
type ResourceRequirement struct {
	// Type 是资源类型
	Type ResourceType
	// Quantity 是资源数量
	Quantity ResourceQuantity
}

// ResourceAllocation 表示资源分配
type ResourceAllocation struct {
	// Type 是资源类型
	Type ResourceType
	// Quantity 是资源数量
	Quantity ResourceQuantity
	// NodeID 是分配资源的节点ID
	NodeID string
}

// ResourceManager 是资源管理器接口
type ResourceManager interface {
	// Start 启动资源管理器
	Start(ctx context.Context) error
	// Stop 停止资源管理器
	Stop(ctx context.Context) error
	// RegisterResource 注册资源
	RegisterResource(ctx context.Context, resourceType ResourceType, quantity ResourceQuantity) error
	// UnregisterResource 注销资源
	UnregisterResource(ctx context.Context, resourceType ResourceType) error
	// AllocateResource 分配资源
	AllocateResource(ctx context.Context, requirement ResourceRequirement) (*ResourceAllocation, error)
	// ReleaseResource 释放资源
	ReleaseResource(ctx context.Context, allocation *ResourceAllocation) error
	// GetAvailableResources 获取可用资源
	GetAvailableResources(ctx context.Context) (map[ResourceType]map[string]ResourceQuantity, error)
	// GetAllocatedResources 获取已分配资源
	GetAllocatedResources(ctx context.Context) (map[ResourceType]map[string]ResourceQuantity, error)
}

// ClusterResourceManager 是基于集群的资源管理器
type ClusterResourceManager struct {
	// manager 是集群管理器
	manager cluster.Manager
	// leaderElection 是领导者选举
	leaderElection coordination.LeaderElection
	// lock 是分布式锁
	lock coordination.DistributedLock
	// localResources 是本地资源
	localResources map[ResourceType]ResourceQuantity
	// allocatedResources 是已分配资源
	allocatedResources map[ResourceType]ResourceQuantity
	// mu 是保护资源的互斥锁
	mu sync.RWMutex
	// logger 是日志记录器
	logger *log.Logger
	// ctx 是上下文
	ctx context.Context
	// cancel 是取消函数
	cancel context.CancelFunc
}

// NewClusterResourceManager 创建一个新的基于集群的资源管理器
func NewClusterResourceManager(manager cluster.Manager) ResourceManager {
	// 创建上下文
	ctx, cancel := context.WithCancel(context.Background())

	// 创建领导者选举
	leaderElection := coordination.NewClusterLeaderElection(manager, "resource-manager")

	// 创建分布式锁
	lock := coordination.NewClusterLock(manager, "resource-manager", 10*time.Second)

	return &ClusterResourceManager{
		manager:            manager,
		leaderElection:     leaderElection,
		lock:               lock,
		localResources:     make(map[ResourceType]ResourceQuantity),
		allocatedResources: make(map[ResourceType]ResourceQuantity),
		logger:             log.New(os.Stderr, "[ResourceManager] ", log.LstdFlags),
		ctx:                ctx,
		cancel:             cancel,
	}
}

// Start 启动资源管理器
func (m *ClusterResourceManager) Start(ctx context.Context) error {
	m.logger.Println("Starting resource manager")

	// 启动领导者选举
	err := m.leaderElection.Start(ctx)
	if err != nil {
		return fmt.Errorf("failed to start leader election: %w", err)
	}

	// 设置领导权变更回调
	m.leaderElection.OnLeadershipChange(func(isLeader bool) {
		if isLeader {
			m.logger.Println("Became resource manager leader")
		} else {
			m.logger.Println("Lost resource manager leadership")
		}
	})

	// 注册本地资源
	err = m.registerLocalResources(ctx)
	if err != nil {
		m.logger.Printf("Failed to register local resources: %v", err)
		// 继续运行，不返回错误
	}

	// 启动定期任务
	go m.periodicTasks()

	return nil
}

// Stop 停止资源管理器
func (m *ClusterResourceManager) Stop(ctx context.Context) error {
	m.logger.Println("Stopping resource manager")

	// 取消上下文
	m.cancel()

	// 停止领导者选举
	err := m.leaderElection.Stop(ctx)
	if err != nil {
		m.logger.Printf("Failed to stop leader election: %v", err)
	}

	return nil
}

// RegisterResource 注册资源
func (m *ClusterResourceManager) RegisterResource(ctx context.Context, resourceType ResourceType, quantity ResourceQuantity) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 更新本地资源
	m.localResources[resourceType] = quantity

	// 更新集群元数据
	err := m.updateResourceMetadata(ctx)
	if err != nil {
		return fmt.Errorf("failed to update resource metadata: %w", err)
	}

	m.logger.Printf("Registered resource: %s = %s", resourceType, quantity.String())
	return nil
}

// UnregisterResource 注销资源
func (m *ClusterResourceManager) UnregisterResource(ctx context.Context, resourceType ResourceType) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 删除本地资源
	delete(m.localResources, resourceType)

	// 更新集群元数据
	err := m.updateResourceMetadata(ctx)
	if err != nil {
		return fmt.Errorf("failed to update resource metadata: %w", err)
	}

	m.logger.Printf("Unregistered resource: %s", resourceType)
	return nil
}

// AllocateResource 分配资源
func (m *ClusterResourceManager) AllocateResource(ctx context.Context, requirement ResourceRequirement) (*ResourceAllocation, error) {
	// 获取锁
	err := m.lock.Lock(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to acquire lock: %w", err)
	}
	defer m.lock.Unlock(ctx)

	// 获取可用资源
	availableResources, err := m.GetAvailableResources(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get available resources: %w", err)
	}

	// 查找满足需求的节点
	var selectedNode string
	var selectedQuantity ResourceQuantity

	for nodeID, resources := range availableResources[requirement.Type] {
		if resources.Value >= requirement.Quantity.Value {
			if selectedNode == "" || resources.Value < selectedQuantity.Value {
				selectedNode = nodeID
				selectedQuantity = resources
			}
		}
	}

	// 如果没有找到满足需求的节点，返回错误
	if selectedNode == "" {
		return nil, fmt.Errorf("no available resources for %s", requirement.Type)
	}

	// 创建资源分配
	allocation := &ResourceAllocation{
		Type:     requirement.Type,
		Quantity: requirement.Quantity,
		NodeID:   selectedNode,
	}

	// 更新已分配资源
	if selectedNode == m.getLocalNodeID() {
		m.mu.Lock()
		defer m.mu.Unlock()

		// 更新本地已分配资源
		current := m.allocatedResources[requirement.Type]
		current.Value += requirement.Quantity.Value
		m.allocatedResources[requirement.Type] = current

		// 更新集群元数据
		err := m.updateResourceMetadata(ctx)
		if err != nil {
			return nil, fmt.Errorf("failed to update resource metadata: %w", err)
		}
	} else {
		// 通知目标节点更新已分配资源
		err := m.notifyResourceAllocation(ctx, selectedNode, allocation)
		if err != nil {
			return nil, fmt.Errorf("failed to notify resource allocation: %w", err)
		}
	}

	m.logger.Printf("Allocated resource: %s = %s on node %s", requirement.Type, requirement.Quantity.String(), selectedNode)
	return allocation, nil
}

// ReleaseResource 释放资源
func (m *ClusterResourceManager) ReleaseResource(ctx context.Context, allocation *ResourceAllocation) error {
	// 获取锁
	err := m.lock.Lock(ctx)
	if err != nil {
		return fmt.Errorf("failed to acquire lock: %w", err)
	}
	defer m.lock.Unlock(ctx)

	// 更新已分配资源
	if allocation.NodeID == m.getLocalNodeID() {
		m.mu.Lock()
		defer m.mu.Unlock()

		// 更新本地已分配资源
		current := m.allocatedResources[allocation.Type]
		current.Value -= allocation.Quantity.Value
		if current.Value < 0 {
			current.Value = 0
		}
		m.allocatedResources[allocation.Type] = current

		// 更新集群元数据
		err := m.updateResourceMetadata(ctx)
		if err != nil {
			return fmt.Errorf("failed to update resource metadata: %w", err)
		}
	} else {
		// 通知目标节点更新已分配资源
		err := m.notifyResourceRelease(ctx, allocation.NodeID, allocation)
		if err != nil {
			return fmt.Errorf("failed to notify resource release: %w", err)
		}
	}

	m.logger.Printf("Released resource: %s = %s on node %s", allocation.Type, allocation.Quantity.String(), allocation.NodeID)
	return nil
}

// GetAvailableResources 获取可用资源
func (m *ClusterResourceManager) GetAvailableResources(ctx context.Context) (map[ResourceType]map[string]ResourceQuantity, error) {
	// 获取所有节点
	nodes, err := m.manager.GetNodes(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get nodes: %w", err)
	}

	// 创建可用资源映射
	availableResources := make(map[ResourceType]map[string]ResourceQuantity)

	// 遍历所有节点
	for _, node := range nodes {
		if node.Status == cluster.NodeStatusUp {
			// 获取节点资源元数据
			resourcesMetadata, err := m.manager.GetMetadata(ctx, node.ID, "resources")
			if err != nil || resourcesMetadata == "" {
				continue
			}

			// 解析资源元数据
			var nodeResources struct {
				Total     map[string]ResourceQuantity `json:"total"`
				Allocated map[string]ResourceQuantity `json:"allocated"`
			}
			if err := json.Unmarshal([]byte(resourcesMetadata), &nodeResources); err != nil {
				m.logger.Printf("Failed to unmarshal resource metadata for node %s: %v", node.ID, err)
				continue
			}

			// 计算可用资源
			for typeStr, total := range nodeResources.Total {
				resourceType := ResourceType(typeStr)
				allocated, ok := nodeResources.Allocated[typeStr]
				if !ok {
					allocated = ResourceQuantity{Value: 0, Unit: total.Unit}
				}

				// 计算可用资源
				available := ResourceQuantity{
					Value: total.Value - allocated.Value,
					Unit:  total.Unit,
				}
				if available.Value < 0 {
					available.Value = 0
				}

				// 添加到可用资源映射
				if _, ok := availableResources[resourceType]; !ok {
					availableResources[resourceType] = make(map[string]ResourceQuantity)
				}
				availableResources[resourceType][node.ID] = available
			}
		}
	}

	return availableResources, nil
}

// GetAllocatedResources 获取已分配资源
func (m *ClusterResourceManager) GetAllocatedResources(ctx context.Context) (map[ResourceType]map[string]ResourceQuantity, error) {
	// 获取所有节点
	nodes, err := m.manager.GetNodes(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get nodes: %w", err)
	}

	// 创建已分配资源映射
	allocatedResources := make(map[ResourceType]map[string]ResourceQuantity)

	// 遍历所有节点
	for _, node := range nodes {
		if node.Status == cluster.NodeStatusUp {
			// 获取节点资源元数据
			resourcesMetadata, err := m.manager.GetMetadata(ctx, node.ID, "resources")
			if err != nil || resourcesMetadata == "" {
				continue
			}

			// 解析资源元数据
			var nodeResources struct {
				Total     map[string]ResourceQuantity `json:"total"`
				Allocated map[string]ResourceQuantity `json:"allocated"`
			}
			if err := json.Unmarshal([]byte(resourcesMetadata), &nodeResources); err != nil {
				m.logger.Printf("Failed to unmarshal resource metadata for node %s: %v", node.ID, err)
				continue
			}

			// 添加到已分配资源映射
			for typeStr, allocated := range nodeResources.Allocated {
				resourceType := ResourceType(typeStr)
				if _, ok := allocatedResources[resourceType]; !ok {
					allocatedResources[resourceType] = make(map[string]ResourceQuantity)
				}
				allocatedResources[resourceType][node.ID] = allocated
			}
		}
	}

	return allocatedResources, nil
}

// registerLocalResources 注册本地资源
func (m *ClusterResourceManager) registerLocalResources(ctx context.Context) error {
	// 注册CPU资源
	cpuCount := runtime.NumCPU()
	err := m.RegisterResource(ctx, ResourceTypeCPU, ResourceQuantity{Value: float64(cpuCount), Unit: "cores"})
	if err != nil {
		return fmt.Errorf("failed to register CPU resource: %w", err)
	}

	// 注册内存资源
	var memStats runtime.MemStats
	runtime.ReadMemStats(&memStats)
	totalMemory := float64(memStats.TotalAlloc) / (1024 * 1024 * 1024) // GB
	err = m.RegisterResource(ctx, ResourceTypeMemory, ResourceQuantity{Value: totalMemory, Unit: "GB"})
	if err != nil {
		return fmt.Errorf("failed to register memory resource: %w", err)
	}

	return nil
}

// updateResourceMetadata 更新资源元数据
func (m *ClusterResourceManager) updateResourceMetadata(ctx context.Context) error {
	// 创建资源元数据
	resourcesMetadata := struct {
		Total     map[string]ResourceQuantity `json:"total"`
		Allocated map[string]ResourceQuantity `json:"allocated"`
	}{
		Total:     make(map[string]ResourceQuantity),
		Allocated: make(map[string]ResourceQuantity),
	}

	// 添加本地资源
	for resourceType, quantity := range m.localResources {
		resourcesMetadata.Total[string(resourceType)] = quantity
	}

	// 添加已分配资源
	for resourceType, quantity := range m.allocatedResources {
		resourcesMetadata.Allocated[string(resourceType)] = quantity
	}

	// 序列化资源元数据
	data, err := json.Marshal(resourcesMetadata)
	if err != nil {
		return fmt.Errorf("failed to marshal resource metadata: %w", err)
	}

	// 更新集群元数据
	err = m.manager.SetMetadata(ctx, "resources", string(data))
	if err != nil {
		return fmt.Errorf("failed to set resource metadata: %w", err)
	}

	return nil
}

// notifyResourceAllocation 通知资源分配
func (m *ClusterResourceManager) notifyResourceAllocation(ctx context.Context, nodeID string, allocation *ResourceAllocation) error {
	// 创建分配消息
	allocationMsg := struct {
		Type     string             `json:"type"`
		Action   string             `json:"action"`
		Resource ResourceAllocation `json:"resource"`
	}{
		Type:     "resource",
		Action:   "allocate",
		Resource: *allocation,
	}

	// 序列化消息
	data, err := json.Marshal(allocationMsg)
	if err != nil {
		return fmt.Errorf("failed to marshal allocation message: %w", err)
	}

	// 发送消息
	err = m.manager.Send(ctx, nodeID, data)
	if err != nil {
		return fmt.Errorf("failed to send allocation message: %w", err)
	}

	return nil
}

// notifyResourceRelease 通知资源释放
func (m *ClusterResourceManager) notifyResourceRelease(ctx context.Context, nodeID string, allocation *ResourceAllocation) error {
	// 创建释放消息
	releaseMsg := struct {
		Type     string             `json:"type"`
		Action   string             `json:"action"`
		Resource ResourceAllocation `json:"resource"`
	}{
		Type:     "resource",
		Action:   "release",
		Resource: *allocation,
	}

	// 序列化消息
	data, err := json.Marshal(releaseMsg)
	if err != nil {
		return fmt.Errorf("failed to marshal release message: %w", err)
	}

	// 发送消息
	err = m.manager.Send(ctx, nodeID, data)
	if err != nil {
		return fmt.Errorf("failed to send release message: %w", err)
	}

	return nil
}

// getLocalNodeID 获取本地节点ID
func (m *ClusterResourceManager) getLocalNodeID() string {
	// 获取本地节点
	localNode, err := m.manager.GetLocalNode(m.ctx)
	if err != nil {
		m.logger.Printf("Failed to get local node: %v", err)
		return ""
	}

	return localNode.ID
}

// periodicTasks 执行定期任务
func (m *ClusterResourceManager) periodicTasks() {
	// 创建定时器
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-m.ctx.Done():
			return
		case <-ticker.C:
			// 更新资源元数据
			err := m.updateResourceMetadata(m.ctx)
			if err != nil {
				m.logger.Printf("Failed to update resource metadata: %v", err)
			}
		}
	}
}
